import akshare as ak
import psycopg2
from psycopg2 import sql
import psycopg2.extras
import logging
import pandas as pd
from datetime import datetime

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def get_db_connection():
    """Establish connection to PostgreSQL database"""
    try:
        conn = psycopg2.connect(
            host="localhost",
            database="stock_db",
            user="leizhen",
            password="751982leizhen",
            port="5432"
        )
        return conn
    except Exception as e:
        logger.error(f"Database connection failed: {e}")
        raise

def create_table_if_not_exists(conn):
    """Create table if it doesn't exist"""
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS stock_gddh_em_df (
        代码 VARCHAR(10),
        简称 VARCHAR(20),
        股东大会名称 VARCHAR(100),
        召开开始日 DATE,
        股权登记日 DATE,
        现场登记日 DATE,
        网络投票时间_开始日 DATE,
        网络投票时间_结束日 DATE,
        决议公告日 DATE,
        公告日 DATE,
        序列号 BIGINT PRIMARY KEY,
        提案 TEXT
    );
    """
    try:
        with conn.cursor() as cursor:
            cursor.execute(create_table_sql)
            conn.commit()
            logger.info("Table created or already exists")
    except Exception as e:
        logger.error(f"Error creating table: {e}")
        conn.rollback()
        raise

def upsert_data(conn, df):
    """Upsert data into the table"""
    upsert_sql = """
    INSERT INTO stock_gddh_em_df (
        代码, 简称, 股东大会名称, 召开开始日, 股权登记日, 现场登记日,
        网络投票时间_开始日, 网络投票时间_结束日, 决议公告日, 公告日, 序列号, 提案
    ) VALUES %s
    ON CONFLICT (序列号) DO UPDATE SET
        代码 = EXCLUDED.代码,
        简称 = EXCLUDED.简称,
        股东大会名称 = EXCLUDED.股东大会名称,
        召开开始日 = EXCLUDED.召开开始日,
        股权登记日 = EXCLUDED.股权登记日,
        现场登记日 = EXCLUDED.现场登记日,
        网络投票时间_开始日 = EXCLUDED.网络投票时间_开始日,
        网络投票时间_结束日 = EXCLUDED.网络投票时间_结束日,
        决议公告日 = EXCLUDED.决议公告日,
        公告日 = EXCLUDED.公告日,
        提案 = EXCLUDED.提案;
    """
    
    try:
        with conn.cursor() as cursor:
            # Convert dataframe to list of tuples, replacing NaT with None
            data = []
            for row in df.itertuples(index=False):
                row_data = []
                for value in row:
                    if pd.isna(value):  # Handles both np.nan and pd.NaT
                        row_data.append(None)
                    else:
                        row_data.append(value)
                data.append(tuple(row_data))
            # Execute the upsert
            psycopg2.extras.execute_values(cursor, upsert_sql, data)
            conn.commit()
            logger.info(f"Upserted {len(data)} records")
    except Exception as e:
        logger.error(f"Error upserting data: {e}")
        conn.rollback()
        raise

def main():
    try:
        # Get data from AKShare and drop duplicates
        logger.info("Fetching data from AKShare...")
        df = ak.stock_gddh_em()
        # Drop duplicates based on sequence number
        df = df.drop_duplicates(subset=['序列号'], keep='first')
        logger.info(f"Found {len(df)} unique records after deduplication")
        
        # Connect to database
        conn = get_db_connection()
        
        # Ensure table exists
        create_table_if_not_exists(conn)
        
        # Upsert data
        upsert_data(conn, df)
        
        logger.info("Data update completed successfully")
        
    except Exception as e:
        logger.error(f"Error in main execution: {e}")
    finally:
        if 'conn' in locals():
            conn.close()

if __name__ == "__main__":
    main()
